
import os, bpy, bpy.ops, math, mathutils
from mathutils import Vector
from rna_xml import rna2xml
import xml.etree.ElementTree as ET
from bpy.types import Operator

def findSorted(xbone, tag):
    all = xbone.findall(tag)
    return all
    if (len(all) == 0):
        return []
    alls = sorted(all, key=lambda bone: bone.get("name"))
    alls.remove(all[0])
    alls = [all[0]] + alls
    return alls

def rnd(v):
    x = round(v.x * 1000, 0)
    y = round(v.y * 1000, 0)
    z = round(v.z * 1000, 0)
    return Vector((x, y, z))

def eqs(v1, v2):
    return (v1.x == v2.x) and (v1.y == v2.y) and (v1.z == v2.z)

def str2vec(s):
    v = Vector((float(i) for i in s.split(' ')))
    if rotZneg90:
        xo = v.x
        v.x = v.y
        v.y = -xo
    return v

def adjustBone(bone, origin, xbone, parent = None):
    pos = str2vec(xbone.get("pos"))
    rot = str2vec(xbone.get("rot"))
    scale = str2vec(xbone.get("scale"))
    
    end = str2vec(xbone.get("end"))
    
    connected[xbone.get("name")] = (xbone.get("connected") == "true")
    
    if (xbone.tag == "bone"):
        piv = str2vec(xbone.get("pivot"))
        p1 = rnd(pos)
        p2 = rnd(piv)
        if not eqs(p1,p2):
            dif = p2-p1
            print("pivot != pos : " + xbone.get("name") + ", diff=" + str(dif))
    
    if not eqs(Vector((0,0,0)), rnd(rot)):
        print("rot != 0 : " + xbone.get("name") + " " + str(rot))
        
    ph = origin + pos
    pt = ph + end
    bone.head = ph
    bone.head_radius = 0.025
    bone.tail = pt
    bone.tail_radius = 0.025
    bone.envelope_weight = 1
    bone.envelope_distance = 0.05
    
    if (parent != None):
        bone.parent = parent

        if connected[bone.name]:
            if not eqs(parent.tail, bone.head):
                print(parent.name + ".tail != " + bone.name + ".head : diff = " + str(bone.head-parent.tail))

    bone.roll = 0

    groupedBones[bone.name] = xbone.get("group")
    
    def addProp(bone, name, val):
        bone[name] = val
    
    if (xbone.tag == "collision_volume"):
        addProp(bone, "restpose_rot_x", rot[0])
        addProp(bone, "restpose_rot_y", rot[1])
        addProp(bone, "restpose_rot_z", rot[2])
        addProp(bone, "restpose_scale_x", scale[0])
        addProp(bone, "restpose_scale_y", scale[0])
        addProp(bone, "restpose_scale_z", scale[0])

    for xb in findSorted(xbone, "bone"):
        createBones(bone, xb)

    for xb in findSorted(xbone, "collision_volume"):
        createBones(bone, xb)

def createBones(parent, xbone):
    boneName = xbone.get("name")
    bone = arm_data.edit_bones.new(boneName)
    origin = parent.head if (parent != None) else Vector((0,0,0))
    adjustBone(bone, origin, xbone, parent)

def adjustBoneGroups():
    for k in groupedBones:
        pose.bones[k].bone_group = pose.bone_groups[groupedBones[k]]
    for bg in boneGroups:
        g = pose.bone_groups[bg[1]]
        i = 1 if len(bg) < 3 else 2
        g.name = bg[0] + " " + bg[i]

def isSelfOrChildOf(bone, name):
    while (bone is not None):
        if (bone.name == name):
            return True
        bone = bone.parent
    return False

def adjustBoneLayers():
    for bone in arm_data.bones:
        for i in range(1, len(bone.layers)):
            bone.layers[i] = False
        pb = pose.bones[bone.name]
        bg = pb.bone_group
        ix = pose.bone_groups.find(bg.name) + 1
        bone.layers[ix] = True
        b = False
        b = b or isSelfOrChildOf(bone, "mCollarLeft")
        b = b or isSelfOrChildOf(bone, "mCollarRight")
        b = b or isSelfOrChildOf(bone, "mHipLeft")
        b = b or isSelfOrChildOf(bone, "mHipRight")
        bone.layers[31] = b
    arm_data.bones["mPelvis"].layers[31] = True
    arm_data.bones["mTorso"].layers[31] = True
    arm_data.bones["mChest"].layers[31] = True
    arm_data.bones["mNeck"].layers[31] = True
    arm_data.bones["mFaceJaw"].layers[31] = True
    arm_data.bones["mFaceRoot"].layers[31] = True

def checkChildrenSymmetry(bone):
    for b in bone.children:
        o = b.name.find("Left");
        if o != -1:
            oppname = b.name.replace("Left", "Right")
            oppbone = arm_data.edit_bones[oppname]
            loc = b.head
            opploc = oppbone.head
            opploc.x = -opploc.x
            if (loc != opploc):
                print("not symmetric (" + b.name.replace("Left", "") + " L/R): " + str(b.head) + " vs. " + str(oppbone.head))
            opploc.x = -loc.x
        checkChildrenSymmetry(b)

def checkSymmetry(s):
    print("Check symmetry" + s)
    checkChildrenSymmetry(arm_data.edit_bones["mPelvis"])
    print("Symmetry checked" + s)

def selectBone(name):
    bone = arm_data.bones[name]
    bone.select = True
    bpy.context.object.data.bones.active = bone
    return bone

def selectPoseBone(name):
    selectBone(name)
    return pose.bones[name]

def selectBoneWithParents(name):
    bone = selectBone(name)
    pb = bone.parent
    while pb != None:
        pb.select = True
        pb = pb.parent

def unselectBones():
    bpy.ops.pose.select_all(action='DESELECT')

def deleteArmatures():
    while (len(bpy.data.armatures) > 0):
        bpy.data.armatures.remove(bpy.data.armatures[0])
    print("orphans-a: " + str(bpy.data.orphans_purge()))

def setMode(mode):
    if (bpy.context.mode != mode):
        bpy.ops.object.mode_set(mode=mode)

def newScene(name):
    setMode("OBJECT")
    scene = bpy.data.scenes.new(name=name)
    bpy.context.window.scene = scene
    return scene

def newArmature():
    setMode("OBJECT")
    bpy.ops.object.armature_add()
    arm_obj = bpy.context.view_layer.objects.active
    arm_data = arm_obj.data
    pose = arm_obj.pose
    setMode("EDIT")
    arm_data.edit_bones.remove(arm_data.edit_bones[0])
    print("orphans-b: " + str(bpy.data.orphans_purge()))
    return arm_obj, arm_data, pose

def createBoneGroups():
    setMode("EDIT")
    for bg in boneGroups:
        pose.bone_groups.new(name=bg[1])
        boneGroup = pose.bone_groups[bg[1]]
        boneGroup.color_set = "THEME" + bg[0]

def loadFileXML(name):
    import ossl.resources as res
    return ET.parse(eval("res." + name.replace(".", "__") + "()"))

def createSkeleton():
    setMode("EDIT")

    createBoneGroups()
    createBones(None, skel.find("bone[@name='mPelvis']"))
    
    checkSymmetry("")
    checkSymmetry(" after adjustment")

    setMode("OBJECT") # use edit_bones as armature.bones 

    print("*** createBones done ***\n")
    createPose()

def createInverseKinematics(dir, angle, tbone_name, pbone_name):
    setMode("EDIT")

    tbone = arm_data.edit_bones[tbone_name]
    pbone = arm_data.edit_bones[pbone_name]
    tbone_parent_name = tbone.parent.name
 
    tikbone = arm_data.edit_bones.new("IK_" + tbone_name)
    ph = tbone.head
    pt = tbone.head + Vector((0,0.1,0))
    tikbone.head = ph
    tikbone.head_radius = 0.025
    tikbone.tail = pt
    tikbone.tail_radius = 0.025
    tikbone.envelope_weight = 1
    tikbone.envelope_distance = 0.05
    tikbone.roll = 0
    tikbone.use_deform = False
    tikbone.layers[31] = True

    pikbone = arm_data.edit_bones.new("IK_" + pbone_name)
    ph = pbone.head - Vector((0,0.3 * dir,0))
    pt = pbone.head - Vector((0,0.4 * dir,0))
    pikbone.head = ph
    pikbone.head_radius = 0.025
    pikbone.tail = pt
    pikbone.tail_radius = 0.025
    pikbone.envelope_weight = 1
    pikbone.envelope_distance = 0.05
    pikbone.roll = 0
    pikbone.use_deform = False
    pikbone.layers[31] = True

    #pbone.head.y += (-dir / 10)
    #pbone.parent.tail.y += (-dir / 10)
    
    n = 1
    p = tbone
    while not p == pbone:
        n += 1
        p = p.parent
        
    setMode("POSE")
    unselectBones()
    selectBone(tbone.parent.name)

    bpy.ops.pose.constraint_add(type="IK")
    
    poseBone = pose.bones[tbone_parent_name]
    
    c = poseBone.constraints[-1]
    c.chain_count = n
    
    c.target = arm_obj
    c.subtarget = tikbone.name

    c.pole_angle = math.radians(angle)
    c.pole_target = arm_obj
    c.pole_subtarget = pikbone.name

    setMode("EDIT")
    tbone = arm_data.edit_bones[tbone_name]
    pbone = arm_data.edit_bones[pbone_name]
    print(tbone.name + " roll: " + str(tbone.roll))
    print(pbone.name + " roll: " + str(pbone.roll))
    setMode("POSE")    

def createPose():
    setMode("POSE")
        
    adjustBoneGroups()
    adjustBoneLayers()

#    createInverseKinematics(2, -90, "mAnkleLeft", "mKneeLeft")
#    createInverseKinematics(2, -90, "mAnkleRight", "mKneeRight")
#    createInverseKinematics(-1, 180, "mWristLeft", "mElbowLeft")
#    createInverseKinematics(-1, 180, "mWristRight", "mElbowRight")
#    
#    bpy.ops.pose.armature_apply()
    print(bpy.context.active_object.pose.bones["mHipRight"].matrix)

    print("*** createPose done ***\n")


#    return ET.parse(os.path.join(filePath, name))
#refDir = "C:\\Blender\\Ruth2 and Roth2 Free Mesh Avatars\Reference-master\\"
#viewerDir = "Linden Lab\\httpsbitbucket.orglindenlabviewer-release\\"
#skel = ET.parse(refDir + viewerDir + "avatar_skeleton.xml")
#lad = ET.parse(refDir + viewerDir + "avatar_lad.xml")
#gen = ET.parse(refDir + viewerDir + "genepool.xml")
#app = ET.parse("C:\\_OSC\Viewer Exports\\FS - Developer bpyMenu - Avatar - Save Appearance to XML\\Owl_Eyes_s_0000.xml")


# ------------------------------------------------------------------------


rotZneg90 = True

boneGroups = [
    ("01", "Torso"),
    ("02", "Spine"),
    ("03", "Arms"),
    ("04", "Hand", "Hands"),
    ("05", "Legs"),
    ("06", "Face"),
    ("07", "Nose"),
    ("08", "Ears"),
    ("09", "Eyes"),
    ("10", "Lips"),
    ("11", "Mouth"),
    ("12", "Wing", "Wings"),
    ("13", "Tail"),
    ("14", "Limb", "Hinds"),
    ("15", "Groin"),
    ("16", "Extra", "Extras"),
    ("17", "Collision")
    ]
groupedBones = {}
connected = {}
arm_obj = None
arm_data = None
pose = None
skel = None

def generate():
    global arm_obj, arm_data, pose, skel
    # deleteArmatures()

    scene = newScene("OSSL Armature")

    arm_obj, arm_data, pose = newArmature()

    skel = loadFileXML("avatar_skeleton.xml")
    createSkeleton()

    setMode("OBJECT")
    
    setMode("POSE")

    def insertRestPose():
        unselectBones()
        selectPoseBone("mPelvis")
        selectPoseBone("PELVIS")
        selectPoseBone("mNeck")
        selectPoseBone("NECK")
        selectPoseBone("mFaceRoot")
        bpy.ops.anim.keyframe_insert_menu(type='Location')

        unselectBones()
        selectPoseBone("mPelvis")
        selectPoseBone("PELVIS")
        selectPoseBone("mNeck")
        selectPoseBone("NECK")
        selectPoseBone("mShoulderLeft")
        selectPoseBone("mHipRight")
        bpy.ops.anim.keyframe_insert_menu(type='Rotation')

    bpy.context.scene.frame_set(0)
    insertRestPose()

    bpy.context.scene.frame_set(12)
    insertRestPose()

    bpy.context.scene.frame_set(84)
    insertRestPose()

    bpy.context.scene.frame_set(96)
    insertRestPose()

    def insertLoopPose(dir):
        
        # arm_data.bone.location is relative to arm_data.parentBone.tail
        unselectBones()
        mPelvis = selectPoseBone("mPelvis")
        print(("mPelvis", arm_data.bones["mPelvis"].head, mPelvis.location))
        mPelvis.location = Vector((dir*0.50,0.30,0.70))
        mPelvis.rotation_quaternion = mathutils.Euler((0,dir*math.pi/8,0)).to_quaternion()
        bpy.ops.anim.keyframe_insert_menu(type='Location')
        bpy.ops.anim.keyframe_insert_menu(type='Rotation')

        unselectBones()
        mFaceRoot = selectPoseBone("mFaceRoot")
        print(("mFaceRoot", arm_data.bones["mFaceRoot"].head, mFaceRoot.location))
        mFaceRoot.location = Vector((0,0.30,0))
        bpy.ops.anim.keyframe_insert_menu(type='Location')

        unselectBones()
        mNeck = selectPoseBone("mNeck")
        print(("mNeck", arm_data.bones["mNeck"].head, mNeck.location))
        mNeck.location = Vector((0,0.60,0))
        mNeck.rotation_quaternion = mathutils.Euler((0,dir*math.pi/2,0)).to_quaternion()
        bpy.ops.anim.keyframe_insert_menu(type='Location')
        bpy.ops.anim.keyframe_insert_menu(type='Rotation')

        unselectBones()
        mShoulderLeft = selectPoseBone("mShoulderLeft")
        print(("mShoulderLeft", arm_data.bones["mShoulderLeft"].head, mShoulderLeft.location))
        mShoulderLeft.rotation_quaternion = mathutils.Euler((math.pi/2,0,0)).to_quaternion()
        bpy.ops.anim.keyframe_insert_menu(type='Rotation')

        unselectBones()
        mHipRight = selectPoseBone("mHipRight")
        print(("mHipRight", arm_data.bones["mHipRight"].head, mHipRight.location))
        mHipRight.rotation_quaternion = mathutils.Euler((math.pi/2,0,0)).to_quaternion()
        bpy.ops.anim.keyframe_insert_menu(type='Rotation')
        print(("mKneeRight", arm_data.bones["mKneeRight"].head, selectPoseBone("mKneeRight").location))

    bpy.context.scene.frame_set(36)
    insertLoopPose(1)

    bpy.context.scene.frame_set(60)
    insertLoopPose(-1)

    bpy.context.scene.frame_set(0)
    scene.frame_start = 0
    scene.frame_end = 96

    #======================================
    
    arm_obj = bpy.context.active_object
    arm_obj = bpy.data.objects[arm_obj.name]
    arm_data = arm_obj.data
    abones = arm_data.bones
    pose = arm_obj.pose
    pbones = pose.bones
    pbgroups = pose.bone_groups
    anim_data = arm_obj.animation_data
    pbone = pbones["mPelvis"]
    abone = abones["mPelvis"]
    anim_data = arm_obj.animation_data
    action = anim_data.action
    scene = bpy.context.scene

    #curves = action.fcurves
    #curve = curves[8] # 
    #frames = curve.keyframe_points
    #kfp = frames[0]


class OSSL_OT_Armature_generate_Default(bpy.types.Operator):
    """Generate Default OSSL Armature with Test Animation"""
    bl_label = "OSSL Menu Armature Generate Default"
    bl_idname = "ossl.armature_generate_default"
    bl_options = { 'REGISTER' }

    def draw(self, context):
        layout = self.layout

    @classmethod
    def poll(cls, context):
        return True
    
    def execute(self, context):
        generate()
        print(arm_data)
        return {'FINISHED'}

def OSSL_operator_Armature_generate_Default(self, comtext):
    self.layout.operator(OSSL_OT_Armature_generate_Default.bl_idname, text="default", icon='OUTLINER_DATA_ARMATURE')

import ossl.util as util

def register():
    print("registering: OSSL Type Armature Generate Default")
    bpy.utils.register_class(OSSL_OT_Armature_generate_Default)
    bpy.types.OSSL_MT_Armature_generate.append(OSSL_operator_Armature_generate_Default)
    print("registered: OSSL Type Armature Generate Default")

def unregister():
    print("unregistering: OSSL Type Armature Generate Default")
    bpy.types.OSSL_MT_Armature_generate.remove(OSSL_operator_Armature_generate_Default)
    bpy.utils.unregister_class(OSSL_OT_Armature_generate_Default)
    print("unregistered: OSSL Type Armature Generate Default")

if __name__ == "__main__":
    import ossl.util as util
    util.reset_addon()
